#
from typing import Dict, List
import logging
import numpy as np
import matplotlib.pyplot as plt
import pymilvus
#
from common.app_registry import AppRegistry as AR
from common.app_manager import AppManager
from mmlm.ce.context_rag import Document, generate_embedding, text_to_chunks, extract_document_batch_embeddings, similarity_search, create_vdb_index, vdb_similarity_search
class ContextRagApp(object):
    def __init__(self):
        self.name = 'apps.ces.context_rag_app.ContextRagApp'

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'RAG增强上下文......v0.0.1')
        if not AR.is_initialized():
            AR.initialize()
        if params['task'] == 1:
            ContextRagApp.embed_doc(params=params)
        elif params['task'] == 2:
            ContextRagApp.simple_search()
        elif params['task'] == 3:
            ContextRagApp.simple_vdb_index(params=params)
        elif params['task'] == 4:
            ContextRagApp.simple_vdb_search(params=params)
        elif params['task'] == 100:
            # 文本嵌入
            ContextRagApp.e001()
        elif params['task'] == 101:
            # 
            ContextRagApp.e002()
        elif params['task'] == 102:
            #
            ContextRagApp.e003()
        else:
            ContextRagApp.t001(params=params)

    @staticmethod
    def embed_doc(params:Dict = {}) -> List[Document]:
        txt_fn = 'docs/radar/chp001.md'
        with open(txt_fn, 'r', encoding='utf-8') as rfd:
            raw_text = rfd.read()
        # 文本分块
        docs = text_to_chunks(raw_text)
        # 求出文本块向量表示
        emb_docs = extract_document_batch_embeddings(docs)
        # print(f'emb_docs: {len(emb_docs)};')
        # print(f'emb: {emb_docs[0].embedding};')
        # print(f'emb_src: {emb_docs[0].content};')
        return emb_docs
    
    @staticmethod
    def simple_search(params:Dict = {}) -> None:
        emb_docs = ContextRagApp.embed_doc()
        query = '非相参累积'
        query_emb = generate_embedding(query)
        query_emb = query_emb.detach().numpy()
        print(f'query_emb: {type(query_emb)};')
        for emb_doc in emb_docs:
            emb_doc.embedding = emb_doc.embedding.detach().cpu().numpy()
        print(f'docs: {type(emb_docs[0].embedding)};')
        rsts = similarity_search(query_embedding=query_emb, documents=emb_docs, top_k=3)
        for doc,sim in rsts:
            print(f'{sim}: {doc.content};')
            print(f'\n********************\n\n\n')

    @staticmethod
    def simple_vdb_index(params:Dict = {}) -> None:
        emb_docs = ContextRagApp.embed_doc()
        query = '非相参累积'
        query_emb = generate_embedding(query)
        query_emb = query_emb.detach().numpy()
        print(f'query_emb: {type(query_emb)};')
        for emb_doc in emb_docs:
            emb_doc.embedding = emb_doc.embedding.detach().cpu().numpy()
        print(f'docs: {type(emb_docs[0].embedding)};')
        create_vdb_index(emb_docs)

    @staticmethod
    def simple_vdb_search(params:Dict = {}) -> None:
        query = '非相参累积'
        query_emb = generate_embedding(query)
        query_emb = query_emb.detach().numpy()
        results = vdb_similarity_search(query_emb.tolist())









    @staticmethod
    def t001(params:Dict = {}) -> None:
        with open('docs/radar/chp001.md', 'r', encoding='utf-8') as rfd:
            raw_text = rfd.read()
        docs = text_to_chunks(raw_text)
        print(f'docs: {len(docs)}; ??????')
        emb_docs = extract_document_batch_embeddings(docs)
        print(f'emb_docs: {len(emb_docs)};')
        print(f'emb: {emb_docs[0].embedding}')
    
    @staticmethod
    def e001() -> None:
        emb_rst = generate_embedding('RAG增强上下文......v0.0.1')
        print(f'emb: {type(emb_rst)}; \n{emb_rst.shape};')

    @staticmethod
    def e002() -> None:
        with open('docs/radar/chp001.md', 'r', encoding='utf-8') as rfd:
            raw_text = rfd.read()
        docs = text_to_chunks(raw_text)
        print(f'docs: {len(docs)}; ??????')

    @staticmethod
    def e003() -> None:
        with open('docs/radar/chp001.md', 'r', encoding='utf-8') as rfd:
            raw_text = rfd.read()
        docs = text_to_chunks(raw_text)
        print(f'docs: {len(docs)}; ??????')
        emb_docs = extract_document_batch_embeddings(docs)
        print(f'emb_docs: {len(emb_docs)};')
        print(f'emb: {emb_docs[0].embedding};')
        print(f'emb_src: {emb_docs[0].content};')